#!/usr/bin/env python3
"""
测试CSVHeaderManager类的功能
"""

import sys
import os
import csv
import tempfile

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from src.plugins.data_analysis.data_analysis import CSVHeaderManager

def create_test_csv_without_header():
    """创建一个没有列名的测试CSV文件"""
    with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False, encoding='utf-8') as f:
        writer = csv.writer(f)
        # 写入数据行（没有列名）
        writer.writerow(['1', 'Alice', '25', 'Engineer'])
        writer.writerow(['2', 'Bob', '30', 'Teacher'])
        writer.writerow(['3', 'Charlie', '35', 'Doctor'])
        return f.name

def create_test_csv_with_header():
    """创建一个有列名的测试CSV文件"""
    with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False, encoding='utf-8') as f:
        writer = csv.writer(f)
        # 写入列名
        writer.writerow(['ID', 'Name', 'Age', 'Occupation'])
        # 写入数据行
        writer.writerow(['1', 'Alice', '25', 'Engineer'])
        writer.writerow(['2', 'Bob', '30', 'Teacher'])
        writer.writerow(['3', 'Charlie', '35', 'Doctor'])
        return f.name

def test_csv_header_manager():
    """测试CSVHeaderManager类"""
    print("=== 测试CSVHeaderManager类 ===")
    
    # 测试1: 为没有列名的CSV文件添加指定列名
    print("\n测试1: 为没有列名的CSV文件添加指定列名")
    test_file1 = create_test_csv_without_header()
    
    plugin = CSVHeaderManager()
    plugin.config = {
        "config": {
            "column_names": ["ID", "Name", "Age", "Occupation"],
            "output_file": "test_output1.csv"
        }
    }
    
    result = plugin.execute(test_file1)
    print(f"结果: {result}")
    
    if result['success']:
        print("生成的CSV文件内容:")
        with open(result['output_file'], 'r', encoding='utf-8') as f:
            print(f.read())
        os.remove(result['output_file'])
    
    os.remove(test_file1)
    
    # 测试2: 自动生成列名
    print("\n测试2: 自动生成列名")
    test_file2 = create_test_csv_without_header()
    
    plugin.config = {
        "config": {
            "auto_generate": True,
            "prefix": "column_",
            "output_file": "test_output2.csv"
        }
    }
    
    result = plugin.execute(test_file2)
    print(f"结果: {result}")
    
    if result['success']:
        print("生成的CSV文件内容:")
        with open(result['output_file'], 'r', encoding='utf-8') as f:
            print(f.read())
        os.remove(result['output_file'])
    
    os.remove(test_file2)
    
    # 测试3: 处理已有列名的CSV文件（has_header=True）
    print("\n测试3: 处理已有列名的CSV文件")
    test_file3 = create_test_csv_with_header()
    
    plugin.config = {
        "config": {
            "column_names": ["编号", "姓名", "年龄", "职业"],
            "has_header": True,
            "output_file": "test_output3.csv"
        }
    }
    
    result = plugin.execute(test_file3)
    print(f"结果: {result}")
    
    if result['success']:
        print("生成的CSV文件内容:")
        with open(result['output_file'], 'r', encoding='utf-8') as f:
            print(f.read())
        os.remove(result['output_file'])
    
    os.remove(test_file3)
    
    # 测试4: 列名数量不匹配的情况
    print("\n测试4: 列名数量不匹配的情况")
    test_file4 = create_test_csv_without_header()
    
    plugin.config = {
        "config": {
            "column_names": ["ID", "Name"],  # 只提供2个列名，但数据有4列
            "auto_generate": True,  # 启用自动生成，补充缺失的列名
            "output_file": "test_output4.csv"
        }
    }
    
    result = plugin.execute(test_file4)
    print(f"结果: {result}")
    
    if result['success']:
        print("生成的CSV文件内容:")
        with open(result['output_file'], 'r', encoding='utf-8') as f:
            print(f.read())
        os.remove(result['output_file'])
    
    os.remove(test_file4)
    
    # 测试5: 配置验证
    print("\n测试5: 配置验证")
    plugin.config = {
        "config": {
            # 没有提供column_names且auto_generate为False
        }
    }
    
    is_valid = plugin.validate_config()
    print(f"配置验证结果: {is_valid}")
    
    print("\n=== 测试完成 ===")

if __name__ == "__main__":
    test_csv_header_manager()